from .targets import Targets
from .scan import Scan
from .scanning_profiles import ScanningProfile
from .vulnerabilities import Vulnerabilities
from export.export import export_to_excel
from config.settings import AWVS_URL, AWVS_API
from typing import Union
import pandas as pd
from utils.log import logger


class AWVSScan(Targets, Scan, ScanningProfile, Vulnerabilities):
    def workflow_targets(self, targets: list):
        """将存活添加到扫描器中 """
        return self.add_targets(targets)

    def workflow_vuln_details(self, target_id: Union[str, list], save_path, sheet_name='vulns'):
        if isinstance(target_id, list):
            results = []
            for target in target_id:
                results.append(self.export_vulns(target))
            vulns = pd.concat(results)
        else:
            vulns = self.export_vulns(target_id)

        self.export(vulns, is_save=True,
                    sheet_name=sheet_name, save_path=save_path)

    def export(self, results, is_save, save_path, sheet_name='AWVS(target)'):
        if not is_save:
            return
        export_to_excel(results, sheet_name=sheet_name,
                        xlsx_name='awvs', save_path=save_path)

    def run(self, alive_domains, is_save=True, save_path=None):
        results = self.workflow_targets(alive_domains)
        export_results = [{'url': x['address'],
                           'target_id': x['target_id']} for x in results]
        self.export_results = export_results
        if results:
            self.export(export_results, is_save, save_path)
        else:
            logger.warning('检索到当前存活域名为空')


class AWVS:
    def __init__(self):
        self.awvs = AWVSScan(AWVS_URL, AWVS_API)

    def run(self, alive_domains, is_save=True, save_path=None):
        self.awvs.run(alive_domains, is_save, save_path)
